#!/usr/bin/env python3
"""
Redis MCP Server - 精简版本
只包含最常用的Redis命令
"""
import json
import logging
from typing import Optional

# 官方MCP SDK导入
from mcp.server.fastmcp import FastMCP

# Redis相关导入
import redis
from redis.exceptions import RedisError

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建MCP服务器实例
mcp = FastMCP("Redis MCP Server - Essential")

# Redis连接配置
REDIS_CONFIG = {
    "host": "localhost",
    "port": 6379,
    "db": 0,
    "decode_responses": True
}

# 全局Redis连接
redis_client: Optional[redis.Redis] = None


def get_redis_client() -> redis.Redis:
    """获取Redis客户端连接"""
    global redis_client
    if redis_client is None:
        redis_client = redis.Redis(**REDIS_CONFIG)
        try:
            redis_client.ping()
            logger.info("Redis连接成功")
        except RedisError as e:
            logger.warning(f"Redis连接失败: {e}")
    return redis_client


# ==================== 字符串操作（最常用） ====================

@mcp.tool()
def redis_set(key: str, value: str, ex: Optional[int] = None) -> str:
    """设置Redis键值对
    
    Args:
        key: Redis键名
        value: 要设置的值
        ex: 过期时间（秒），可选
    """
    try:
        client = get_redis_client()
        client.set(key, value, ex=ex)
        if ex:
            return f"✅ 成功设置键 '{key}' = '{value}'，{ex}秒后过期"
        else:
            return f"✅ 成功设置键 '{key}' = '{value}'"
    except RedisError as e:
        return f"❌ 设置失败: {str(e)}"


@mcp.tool()
def redis_get(key: str) -> str:
    """获取Redis键的值
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        value = client.get(key)
        if value is None:
            return f"❌ 键 '{key}' 不存在"
        return f"✅ 键 '{key}' 的值: {value}"
    except RedisError as e:
        return f"❌ 获取失败: {str(e)}"


@mcp.tool()
def redis_delete(key: str) -> str:
    """删除Redis键
    
    Args:
        key: 要删除的Redis键名
    """
    try:
        client = get_redis_client()
        result = client.delete(key)
        if result:
            return f"✅ 成功删除键 '{key}'"
        else:
            return f"❌ 键 '{key}' 不存在"
    except RedisError as e:
        return f"❌ 删除失败: {str(e)}"


@mcp.tool()
def redis_exists(key: str) -> str:
    """检查Redis键是否存在
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        exists = client.exists(key)
        if exists:
            return f"✅ 键 '{key}' 存在"
        else:
            return f"❌ 键 '{key}' 不存在"
    except RedisError as e:
        return f"❌ 检查失败: {str(e)}"


@mcp.tool()
def redis_keys(pattern: str = "*") -> str:
    """查询匹配模式的所有键
    
    Args:
        pattern: 匹配模式，默认为 "*"（所有键）
    """
    try:
        client = get_redis_client()
        keys = client.keys(pattern)
        if keys:
            return f"✅ 找到 {len(keys)} 个键: {', '.join(keys)}"
        else:
            return f"❌ 没有找到匹配 '{pattern}' 的键"
    except RedisError as e:
        return f"❌ 查询失败: {str(e)}"


@mcp.tool()
def redis_incr(key: str) -> str:
    """将键的值递增1
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        value = client.incr(key)
        return f"✅ 键 '{key}' 递增后的值: {value}"
    except RedisError as e:
        return f"❌ 递增失败: {str(e)}"


@mcp.tool()
def redis_decr(key: str) -> str:
    """将键的值递减1
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        value = client.decr(key)
        return f"✅ 键 '{key}' 递减后的值: {value}"
    except RedisError as e:
        return f"❌ 递减失败: {str(e)}"


# ==================== 列表操作（常用） ====================

@mcp.tool()
def redis_lpush(key: str, value: str) -> str:
    """向列表左侧插入元素
    
    Args:
        key: 列表键名
        value: 要插入的值
    """
    try:
        client = get_redis_client()
        count = client.lpush(key, value)
        return f"✅ 成功向列表 '{key}' 左侧插入元素，列表长度: {count}"
    except RedisError as e:
        return f"❌ 插入失败: {str(e)}"


@mcp.tool()
def redis_rpush(key: str, value: str) -> str:
    """向列表右侧插入元素
    
    Args:
        key: 列表键名
        value: 要插入的值
    """
    try:
        client = get_redis_client()
        count = client.rpush(key, value)
        return f"✅ 成功向列表 '{key}' 右侧插入元素，列表长度: {count}"
    except RedisError as e:
        return f"❌ 插入失败: {str(e)}"


@mcp.tool()
def redis_lpop(key: str) -> str:
    """从列表左侧弹出元素
    
    Args:
        key: 列表键名
    """
    try:
        client = get_redis_client()
        value = client.lpop(key)
        if value is not None:
            return f"✅ 从列表 '{key}' 左侧弹出元素: {value}"
        else:
            return f"❌ 列表 '{key}' 为空或不存在"
    except RedisError as e:
        return f"❌ 弹出失败: {str(e)}"


@mcp.tool()
def redis_rpop(key: str) -> str:
    """从列表右侧弹出元素
    
    Args:
        key: 列表键名
    """
    try:
        client = get_redis_client()
        value = client.rpop(key)
        if value is not None:
            return f"✅ 从列表 '{key}' 右侧弹出元素: {value}"
        else:
            return f"❌ 列表 '{key}' 为空或不存在"
    except RedisError as e:
        return f"❌ 弹出失败: {str(e)}"


@mcp.tool()
def redis_lrange(key: str, start: int = 0, end: int = -1) -> str:
    """获取列表指定范围的元素
    
    Args:
        key: 列表键名
        start: 开始索引，默认0
        end: 结束索引，默认-1（到末尾）
    """
    try:
        client = get_redis_client()
        values = client.lrange(key, start, end)
        if values:
            return f"✅ 列表 '{key}' 的元素: {values}"
        else:
            return f"❌ 列表 '{key}' 为空或不存在"
    except RedisError as e:
        return f"❌ 获取失败: {str(e)}"


@mcp.tool()
def redis_llen(key: str) -> str:
    """获取列表长度
    
    Args:
        key: 列表键名
    """
    try:
        client = get_redis_client()
        length = client.llen(key)
        return f"✅ 列表 '{key}' 的长度: {length}"
    except RedisError as e:
        return f"❌ 获取长度失败: {str(e)}"


# ==================== 哈希操作（常用） ====================

@mcp.tool()
def redis_hset(key: str, field: str, value: str) -> str:
    """设置哈希字段值
    
    Args:
        key: 哈希键名
        field: 字段名
        value: 字段值
    """
    try:
        client = get_redis_client()
        result = client.hset(key, field, value)
        if result:
            return f"✅ 成功设置哈希 '{key}' 的字段 '{field}' = '{value}'"
        else:
            return f"✅ 更新哈希 '{key}' 的字段 '{field}' = '{value}'"
    except RedisError as e:
        return f"❌ 设置失败: {str(e)}"


@mcp.tool()
def redis_hget(key: str, field: str) -> str:
    """获取哈希字段值
    
    Args:
        key: 哈希键名
        field: 字段名
    """
    try:
        client = get_redis_client()
        value = client.hget(key, field)
        if value is None:
            return f"❌ 哈希 '{key}' 的字段 '{field}' 不存在"
        return f"✅ 哈希 '{key}' 的字段 '{field}' 值: {value}"
    except RedisError as e:
        return f"❌ 获取失败: {str(e)}"


@mcp.tool()
def redis_hgetall(key: str) -> str:
    """获取哈希的所有字段和值
    
    Args:
        key: 哈希键名
    """
    try:
        client = get_redis_client()
        data = client.hgetall(key)
        if data:
            formatted = json.dumps(data, ensure_ascii=False, indent=2)
            return f"✅ 哈希 '{key}' 的所有数据:\n{formatted}"
        else:
            return f"❌ 哈希 '{key}' 为空或不存在"
    except RedisError as e:
        return f"❌ 获取失败: {str(e)}"


@mcp.tool()
def redis_hdel(key: str, field: str) -> str:
    """删除哈希字段
    
    Args:
        key: 哈希键名
        field: 要删除的字段名
    """
    try:
        client = get_redis_client()
        count = client.hdel(key, field)
        if count:
            return f"✅ 成功从哈希 '{key}' 删除字段 '{field}'"
        else:
            return f"❌ 哈希 '{key}' 的字段 '{field}' 不存在"
    except RedisError as e:
        return f"❌ 删除失败: {str(e)}"


# ==================== 集合操作（常用） ====================

@mcp.tool()
def redis_sadd(key: str, member: str) -> str:
    """向集合添加成员
    
    Args:
        key: 集合键名
        member: 要添加的成员
    """
    try:
        client = get_redis_client()
        count = client.sadd(key, member)
        if count:
            return f"✅ 成功向集合 '{key}' 添加成员 '{member}'"
        else:
            return f"❌ 成员 '{member}' 已在集合 '{key}' 中"
    except RedisError as e:
        return f"❌ 添加失败: {str(e)}"


@mcp.tool()
def redis_smembers(key: str) -> str:
    """获取集合的所有成员
    
    Args:
        key: 集合键名
    """
    try:
        client = get_redis_client()
        members = client.smembers(key)
        if members:
            return f"✅ 集合 '{key}' 的成员: {list(members)}"
        else:
            return f"❌ 集合 '{key}' 为空或不存在"
    except RedisError as e:
        return f"❌ 获取失败: {str(e)}"


@mcp.tool()
def redis_srem(key: str, member: str) -> str:
    """从集合删除成员
    
    Args:
        key: 集合键名
        member: 要删除的成员
    """
    try:
        client = get_redis_client()
        count = client.srem(key, member)
        if count:
            return f"✅ 成功从集合 '{key}' 删除成员 '{member}'"
        else:
            return f"❌ 成员 '{member}' 不在集合 '{key}' 中"
    except RedisError as e:
        return f"❌ 删除失败: {str(e)}"


# ==================== 通用操作（必需） ====================

@mcp.tool()
def redis_expire(key: str, seconds: int) -> str:
    """设置键的过期时间
    
    Args:
        key: Redis键名
        seconds: 过期时间（秒）
    """
    try:
        client = get_redis_client()
        result = client.expire(key, seconds)
        if result:
            return f"✅ 成功设置键 '{key}' 在 {seconds} 秒后过期"
        else:
            return f"❌ 键 '{key}' 不存在"
    except RedisError as e:
        return f"❌ 设置过期时间失败: {str(e)}"


@mcp.tool()
def redis_ttl(key: str) -> str:
    """获取键的剩余生存时间
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        ttl = client.ttl(key)
        if ttl == -1:
            return f"✅ 键 '{key}' 永不过期"
        elif ttl == -2:
            return f"❌ 键 '{key}' 不存在"
        else:
            return f"✅ 键 '{key}' 剩余生存时间: {ttl} 秒"
    except RedisError as e:
        return f"❌ 获取TTL失败: {str(e)}"


@mcp.tool()
def redis_type(key: str) -> str:
    """获取键的数据类型
    
    Args:
        key: Redis键名
    """
    try:
        client = get_redis_client()
        key_type = client.type(key)
        return f"✅ 键 '{key}' 的类型: {key_type}"
    except RedisError as e:
        return f"❌ 获取类型失败: {str(e)}"


@mcp.tool()
def redis_ping() -> str:
    """测试Redis连接"""
    try:
        client = get_redis_client()
        result = client.ping()
        if result:
            return f"✅ Redis连接正常: PONG"
        else:
            return f"❌ Redis连接异常"
    except RedisError as e:
        return f"❌ 连接测试失败: {str(e)}"


@mcp.tool()
def redis_info() -> str:
    """获取Redis服务器信息"""
    try:
        client = get_redis_client()
        info = client.info()
        
        # 提取关键信息
        key_info = {
            "redis_version": info.get("redis_version", "未知"),
            "used_memory_human": info.get("used_memory_human", "未知"),
            "connected_clients": info.get("connected_clients", 0),
            "total_commands_processed": info.get("total_commands_processed", 0),
            "keyspace_hits": info.get("keyspace_hits", 0),
            "keyspace_misses": info.get("keyspace_misses", 0)
        }
        
        formatted = json.dumps(key_info, ensure_ascii=False, indent=2)
        return f"✅ Redis服务器信息:\n{formatted}"
    except RedisError as e:
        return f"❌ 获取服务器信息失败: {str(e)}"


# 主函数
if __name__ == "__main__":
    # 直接运行服务器
    mcp.run()
